Repository navigation
feat: add opt-in single_host_tpu_init for node-local diagnostics on multi-host slices - #144
Open
simrankaurb wants to merge 1 commit into
Open
simrankaurb wants to merge 1 commit into
simrankaurb wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a helper method _replicated_normal to generate replicated arrays of normal samples locally on each device, preventing device-to-device transfers that can hang when ICI routing is not configured. The generate_inputs method is updated to utilize this helper for input and scaling factor generation. The review feedback suggests several improvements: simplifying the creation of replicated_sharding with an empty PartitionSpec(), using the idiomatic self.mesh.local_devices property instead of manual filtering, and catching both RuntimeError and IndexError when accessing the CPU device to prevent potential crashes.
Comment on lines
+215
to
+217
| replicated_sharding = jax.sharding.NamedSharding( | ||
| self.mesh, jax.sharding.PartitionSpec(*([None] * len(shape))) | ||
| ) |
There was a problem hiding this comment.
We can simplify the creation of replicated_sharding by using an empty PartitionSpec(). In JAX, an empty PartitionSpec() represents a fully replicated sharding across all dimensions of the mesh, which is cleaner and matches the usage elsewhere in this file (e.g., line 524).
Suggested change
| replicated_sharding = jax.sharding.NamedSharding( | |
| self.mesh, jax.sharding.PartitionSpec(*([None] * len(shape))) | |
| ) | |
| replicated_sharding = jax.sharding.NamedSharding( | |
| self.mesh, jax.sharding.PartitionSpec() | |
| ) |
Comment on lines
+225
to
+227
| local_devices = [ | ||
| d for d in self.mesh.devices.flat if d.process_index == jax.process_index() | ||
| ] |
There was a problem hiding this comment.
Instead of manually filtering self.mesh.devices.flat by process_index, we can use the built-in self.mesh.local_devices property. This is more idiomatic, cleaner, and less error-prone.
local_devices = self.mesh.local_devices
Comment on lines
+221
to
+224
| try: | ||
| key = jax.device_put(key, jax.devices("cpu")[0]) | ||
| except RuntimeError: | ||
| pass # No CPU backend (JAX_PLATFORMS=tpu); fall back to device key. |
There was a problem hiding this comment.
If the CPU backend is registered but has no devices, jax.devices("cpu") would return an empty list, and accessing [0] would raise an IndexError. It is safer to catch both RuntimeError and IndexError to ensure robust fallback behavior.
Suggested change
| try: | |
| key = jax.device_put(key, jax.devices("cpu")[0]) | |
| except RuntimeError: | |
| pass # No CPU backend (JAX_PLATFORMS=tpu); fall back to device key. | |
| try: | |
| key = jax.device_put(key, jax.devices("cpu")[0]) | |
| except (RuntimeError, IndexError): | |
| pass # No CPU backend (JAX_PLATFORMS=tpu); fall back to device key. |
Comment on lines
+257
to
+260
| try: | ||
| key = jax.device_put(key, jax.devices("cpu")[0]) | ||
| except RuntimeError: | ||
| pass # No CPU backend (JAX_PLATFORMS=tpu); fall back to device key. |
There was a problem hiding this comment.
Similarly, we should catch both RuntimeError and IndexError here to handle cases where the CPU backend is registered but has no devices.
Suggested change
| try: | |
| key = jax.device_put(key, jax.devices("cpu")[0]) | |
| except RuntimeError: | |
| pass # No CPU backend (JAX_PLATFORMS=tpu); fall back to device key. | |
| try: | |
| key = jax.device_put(key, jax.devices("cpu")[0]) | |
| except (RuntimeError, IndexError): | |
| pass # No CPU backend (JAX_PLATFORMS=tpu); fall back to device key. |
simrankaurb
force-pushed
the
fix/gemm-replicated-inputs-no-d2d
branch
from
October 26, 2026 12:05
d6baf11 to
b78b61c
Compare
simrankaurb
changed the title
…-host slices
Adds a benchmark-level `single_host_tpu_init: true` config key and a
`--single_host_tpu_init` common execution flag (default off, accepted by
both `benchmark run` and `benchmark run-config`) that initialize the
current host standalone inside a multi-host slice, so single-host
benchmarks such as `gemm` and `hbm` can be used for node-local diagnostics
(e.g. Cluster Health Scanner) without the other hosts of the slice
participating.
When enabled, `core/single_host.py` (applied after op_flags and before the
first JAX device call) forces a 1-host topology (`TPU_HOST_BOUNDS`,
`TPU_WORKER_ID`, `TPU_WORKER_HOSTNAMES`, `TPU_SKIP_MDS_QUERY`), skips
SliceBuilder/tpunetd (`--deepsea_hal_test_skip_slicebuilder=true`,
`--deepsea_hal_test_allow_multichip_skip_slicebuilder=true`) and disables
the cross-host launch barrier. LIBTPU_INIT_ARGS is merged by flag name so
pre-existing values are overridden rather than duplicated, and every
overridden env var / flag is logged with its previous value.
Because the slice's ICI routing is left untouched, device-to-device
transfers hang in this mode. `GeneralizedGemmBenchmark.generate_inputs`
previously sampled the operands on the default device and broadcast them
with `jax.device_put(x, replicated_sharding)`, which hung for any matrix
size. It now generates each operand under `jax.jit(...,
out_shardings=replicated_sharding)` (the pattern already used by `hbm.py`),
so every device produces the same replicated values and no broadcast is
needed. This is a single code path for both modes: with the flag off the
only behavioural change is that GEMM operand generation no longer
broadcasts over ICI (values are bit-identical). Only `gemm`,
`gemm_throttling` (same input path) and `hbm` are supported in this mode;
the runner prints a warning when the flag is used with any other benchmark.
Verification:
- TPU, earlier revision of this change (replicated operands built via
`jax.make_array_from_single_device_arrays` instead of
`jax.jit(out_shardings=...)`), on a tpu7x-standard-4t host of a 2x2x4
slice (jax 0.12.0, libtpu 0.0.50), no manual environment set up:
`single_host_tpu_init: true`, gemm 16384^3 bf16 x20, xprof
local_slowest_device: 2127 TFLOP/s per chip, 20/20 xprof runs, 78 s;
`--single_host_tpu_init` gemm 1024^3: 20 s; hbm copy 1G elems:
2790 GB/s; control without the flag on the same node fails at init
("Expected 4 worker addresses, got 1").
- CPU: the current `jax.jit(out_shardings=...)` path yields arrays
bit-identical to the previous default path with fully replicated,
identical shards on a 4-device CPU mesh; it has not yet been re-run on
TPU. New unit tests cover config/CLI plumbing, LIBTPU_INIT_ARGS merging,
runner stage ordering and the unsupported-benchmark warning.
simrankaurb
force-pushed
the
fix/gemm-replicated-inputs-no-d2d
branch
from
October 28, 2026 06:51
b78b61c to
a80aaba
Compare
|
अगर तुम ऐसा AI बनाना चाहते हो जिसकी दुनिया में करोड़ों लोगों को जरूरत हो, तो मेरी सलाह है कि केवल एक और ChatGPT जैसा AI मत बनाओ। ऐसा AI बनाओ जो लोगों की कोई बड़ी, रोज़ की समस्या हल करे और उनका समय, पैसा या मेहनत बचाए।
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds an opt‑in
single_host_tpu_init(YAML keysingle_host_tpu_init: true/ CLI--single_host_tpu_init, default off) that initializes the current host standalone inside a multi‑host slice, so single‑host benchmarks (gemm,hbm, …) can be used for node‑local diagnostics — e.g. by Cluster Health Scanner — without the other hosts of the slice participating. No existing behaviour changes when the flag is off (thematmul.pydiff is pure additions). This brings thechsbranch's single‑host init path (#142) tomain.What the flag does
core/single_host.py::configure_single_host_tpu(), called by the runner afterop_flags.yamlis applied and before the first JAX device call:TPU_HOST_BOUNDS=1,1,1,TPU_WORKER_ID=0,TPU_WORKER_HOSTNAMES=localhost,TPU_SKIP_MDS_QUERY=trueLIBTPU_INIT_ARGS:--deepsea_hal_test_skip_slicebuilder=true --deepsea_hal_test_allow_multichip_skip_slicebuilder=true --xla_tpu_use_enhanced_launch_barrier=false(and flips an existingenhanced_launch_barrier=true)Because the slice's ICI routing is left untouched, device‑to‑device transfers hang in this mode.
GeneralizedGemmBenchmark.generate_inputssamples the operands on the default device and then callsjax.device_put(x, replicated_sharding), which broadcasts over ICI — sogemmhung for any matrix size, with and without xprof (Slow PjRt TPU operation detectedon all local devices). When the flag is on, the operands (and optional scaling factors / accumulator) are instead generated independently on each local mesh device from a CPU‑pinned key and assembled withjax.make_array_from_single_device_arrays; the only transfers are host→device. Same sharding (P(None, None)) and same contents (same key everywhere). A key left on TPU:0 was verified to still hang — the per‑devicedevice_putof the key is itself a D2D copy.Plumbing mirrors
xprof_timing:BenchmarkKey.SINGLE_HOST_TPU_INIT+ bool validation inconfig.py,--single_host_tpu_initonbenchmark run, pop/forward inrun-config,runner.run_benchmarks(single_host_tpu_init=…). README updated.Verification
tpu7x‑standard‑4t host of a 2x2x4 slice, jax 0.12.0 / libtpu 0.0.50, no manual environment set up in the pod:
single_host_tpu_init: true, gemm 16384³ bf16 ×20,xprof_device_timing_aggregation: local_slowest_device--single_host_tpu_initgemm 1024³--single_host_tpu_inithbm copy, 1 G elems, device 0Invalid --deepsea_slice_builder_worker_addresses … Expected 4 worker addresses, got 1CPU (4‑device host mesh): flagged and default input paths produce identical replicated arrays and bit‑identical GEMM output for the same seed (incl. scaling‑factor + accumulator inputs);
single_host_tpu_init: "yes"is rejected with a clearValueError; the CLI flag parses and is not forwarded intoGemmParams.Note: upstream
tests/importaccelerator_microbenchmarks.tests, which is not resolvable in the OSS checkout, so they were not run here.